#include "yolotest.h"
#include <iostream>

/**
 * @brief 构造函数
 */
yolo_detector::yolo_detector(){
    inputSize = cv::Size(640,640);
    class_names = {
        "default", "Class_1", "Class_2", "Class_3",
        "Class_4", "Class_5", "Class_6", "Class_7",
        "Class_8", "Class_9", "Class_10", "Class_11", "Class_12"
    };
}

/**
 * @brief 初始化函数
 */
bool yolo_detector::initialize(const std::string& model_path){
    try{
        // 读取ONNX模型
        net = cv::dnn::readNetFromONNX(model_path);

        //选定内部参数
        net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
        net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
        return true;
        std::cout<< "ONNX文件读取成功" << std::endl;
    }catch(const cv::Exception& e){
        std::cerr<< "ONNX文件读取失败，请检查路径" <<std::endl;
        return false;
    }
}

/**
 * @brief 预处理函数 
 */
cv::Mat yolo_detector::preprocess(const cv::Mat& image){
    // 创建神经网络
    cv::Mat blob;

    // 归一化处理
    cv::dnn::blobFromImage(
        image,              // 输入图像
        blob,               // 输出神经网络
        1.0/255.0,          // 归一化系数
        inputSize,         // 初始化图像尺寸大小
        cv::Scalar(0,0,0),  // mean
        true,               // 是否交换BGR通道
        false,              // 是否输出为单通道图像
        CV_32F              // 输出结果类型
    );

    return blob;
}


/**
 * @brief 目标检测函数
 */
std::vector<result> yolo_detector::detect(const cv::Mat& image){
    // 检查图像是否为空
    if(image.empty()){
        std::cerr<< "图像为空，请检查路径" <<std::endl;
    }

    // 1、预处理
    cv::Mat blob = preprocess(image);

    // 2、加载网络
    net.setInput(blob);

    // 3、获取输出层名称
    std::vector<std::string> OutLayarNames = net.getUnconnectedOutLayersNames();

    // 4、执行向前推理
    std::vector<cv::Mat> Outputs;
    net.forward(Outputs,OutLayarNames);

    // 5、输出基本信息
    std::cout << "输出模型数量：" << Outputs.size() << std::endl;
    if(!Outputs.empty()){
        std::cout << "数据输出为：[";
        for(int i = 0;i < Outputs[0].dims;i++){
            std::cout << Outputs[0].size[i] << ",";
        }
        std::cout << "]" << std::endl;
    }
    
    // 6、后处理
    return postprocess(Outputs[0], image.size(), 0.25f, 0.45f);
}

std::vector<result> yolo_detector::postprocess(const cv::Mat& Output, const cv::Size& OriginalImageSize, float conf = 0.25f, float nms = 0.45f){
    // 定义输出
    std::vector<result> results;

    // 数据处理
    cv::Mat resizeOutput;
    if(Output.dims == 3){
        resizeOutput = Output;
    }else if(Output.dims == 2){
        cv::merge(std::vector<cv::Mat>{Output},resizeOutput);
        std::cout << "输出维度为2,已手动纠正" << std::endl;
    }else if(Output.dims == 4){
        std::vector<int> newShape = {
            Output.size[1],
            Output.size[2],
            Output.size[3]
        };
        resizeOutput = Output.reshape(0,newShape);
        std::cout << "输出维度为4,已手动纠正" << std::endl;
    }
    
    // 计算缩放比例
    float xScalar = static_cast<float>(OriginalImageSize.width) / inputSize.width;
    float yScalar = static_cast<float>(OriginalImageSize.height) / inputSize.height;

    // 储存结果的容器
    std::vector<int> classIds;
    std::vector<float> confidences;
    std::vector<cv::Rect> boxes;
    
    // 解析模型输出
    int num_classes = resizeOutput.size[1] - 4; // 类别数量 = 总行数 - 4
    int num_anchors = resizeOutput.size[2];     // 锚点数量 = 总列数

    std::cout << "检测到锚点数：" << num_anchors << ",检测到类别数：" << num_classes << std::endl;

    // 遍历所有锚点（检测结果）
    for(int col = 0; col < num_anchors; col++){
        // 类别概率表格
        std::vector<float> classScores(num_classes);
        for(int cls = 0; cls < num_classes; cls++){
            classScores[cls] = resizeOutput.ptr<float>(0)[num_anchors * (cls + 4) + col];
        }

        // 找出最高置信度及其对应的类别
        float maxClassScore = 0;
        int classIdPoint = 0;
        for(int cls = 0; cls < num_classes; cls++){
            if(classScores[cls] > maxClassScore){
                maxClassScore = classScores[cls];
                classIdPoint = cls;
            }
        }

        // 筛选高于阈值的检测结果
        if(maxClassScore > conf){
            try{
                // 获取边界框坐标（XYWH格式）
                float x = resizeOutput.ptr<float>(0)[num_anchors * 0 + col];        // 中心点x坐标
                float y = resizeOutput.ptr<float>(0)[num_anchors * 1 + col];        // 中心点y坐标
                float width = resizeOutput.ptr<float>(0)[num_anchors * 2 + col];    // 边界框宽
                float height = resizeOutput.ptr<float>(0)[num_anchors * 3 + col];   // 边界框高

                // 转换为左上角坐标并缩放到原图大小
                int left = static_cast<int>((x - width/2) * xScalar);
                int top = static_cast<int>((y - height/2) * yScalar);
                int original_w = static_cast<int>(width * xScalar);
                int original_h = static_cast<int>(height * yScalar);

                // 保存边界框
                boxes.push_back(cv::Rect(left, top, original_w, original_h));

                // 验证并保存识别id
                if(classIdPoint > 0 && classIdPoint < static_cast<int>(class_names.size())){
                    classIds.push_back(classIdPoint);
                }else{
                    classIds.push_back(0);
                }

                // 保存置信度
                confidences.push_back(maxClassScore); 
            }catch(const cv::Exception& e){
                std::cerr << "处理第" << col << "列时发生错误" << e.what() << std::endl;
            }
        }
    }

    // 应用非极大值抑制(NMS), 消除重叠框
    std::vector<int> indices;
    if(!boxes.empty()){
        cv::dnn::NMSBoxes(boxes, confidences, conf, nms, indices);
    }

    //整合最终结果
    for(size_t i = 0; i < indices.size(); i++){
        int idx = indices[i];
        result Result;
        Result.class_id = classIds[idx];
        Result.class_name = class_names[Result.class_id];
        Result.confidence = confidences[idx];
        Result.box = boxes[idx];
        
        results.push_back(Result);
    }
    return results;
}

void yolo_detector::drawresluts(const cv::Mat& image,std::vector<result> results){
    for(const auto& result : results){
        // 根据id生成边界框颜色
        cv::Scalar color(0, 255 * (result.class_id * 123 % 255) / 255.0, 255 * (result.class_id * 213 % 255) / 255.0);

        // 绘制边界框
        cv::rectangle(image, result.box, color, 2);

        // 绘制标签文本
        std::string label = result.class_name + ":" + std::to_string(static_cast<int>(result.confidence * 100)) + "%";

        // 绘制标签背景
        int baseLine;
        cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
        cv::rectangle(image, 
                     cv::Point(result.box.x, result.box.y - labelSize.height - baseLine),
                     cv::Point(result.box.x + labelSize.width, result.box.y),
                     color, 
                     cv::FILLED);
                     
        // 绘制标签文本
        cv::putText(image, 
                   label, 
                   cv::Point(result.box.x, result.box.y - baseLine),
                   cv::FONT_HERSHEY_SIMPLEX, 
                   0.5, 
                   cv::Scalar(255, 255, 255), 
                   1);
    }
}
